package leetcode

import (
	"container/heap"
	"sort"
)

type MedianFinder struct {
	minQ *IntHeap
	maxQ *IntHeap
}

func Constructor() MedianFinder {
	return MedianFinder{
		minQ: &IntHeap{make(sort.IntSlice, 0)},
		maxQ: &IntHeap{make(sort.IntSlice, 0)},
	}
}

func (f *MedianFinder) AddNum(num int) {
	// 保证 0 <= minQ - maxQ <= 1
	if f.minQ.Len() == 0 || num <= -f.minQ.IntSlice[0] {
		heap.Push(f.minQ, -num)
		if f.maxQ.Len()+1 < f.minQ.Len() {
			heap.Push(f.maxQ, -heap.Pop(f.minQ).(int))
		}
	} else {
		heap.Push(f.maxQ, num)
		if f.maxQ.Len() > f.minQ.Len() {
			heap.Push(f.minQ, -heap.Pop(f.maxQ).(int))
		}
	}
	return
}

func (f *MedianFinder) FindMedian() float64 {
	if f.minQ.Len() > f.maxQ.Len() {
		return float64(-f.minQ.IntSlice[0])
	}
	return float64(f.maxQ.IntSlice[0]-f.minQ.IntSlice[0]) / 2
}

type IntHeap struct {
	sort.IntSlice
}

func (h *IntHeap) Push(x any) {
	h.IntSlice = append(h.IntSlice, x.(int))
}

func (h *IntHeap) Pop() any {
	n := len(h.IntSlice)
	x := h.IntSlice[n-1]
	h.IntSlice = h.IntSlice[:n-1]
	return x
}
